import {
  Box,
  FormControl,
  FormLabel,
  Input,
  Link,
  Select,
  Sheet,
  Table,
  Typography,
  Option,
} from '@mui/joy';
import {
  ArrowDownIcon,
  DownloadCloudIcon,
  FlaskRoundIcon,
  InfoIcon,
  SearchIcon,
  StoreIcon,
  Trash2Icon,
} from 'lucide-react';
import React, { useState } from 'react';

import * as chatAPI from '../../lib/transformerlab-api-sdk';

import { filterByFilters, licenseTypes, modelTypes } from '../../lib/utils';
import TinyMLXLogo from '../Shared/TinyMLXLogo';
import SelectButton from '../Experiment/SelectButton';
import { Link as ReactRouterLink } from 'react-router-dom';

type Order = 'asc' | 'desc';

const LocalModelsTable = ({
  models,
  mutateModels,
  setFoundation,
  setAdaptor,
  pickAModelMode = false,
}) => {
  const [order, setOrder] = useState<Order>('desc');
  const [searchText, setSearchText] = useState('');
  const [filters, setFilters] = useState({});

  const renderFilters = () => (
    <>
      <FormControl size="sm">
        <FormLabel>License</FormLabel>
        <Select
          placeholder="Filter by license"
          slotProps={{ button: { sx: { whiteSpace: 'nowrap' } } }}
          value={filters?.license}
          disabled
          onChange={(e, newValue) => {
            setFilters({ ...filters, license: newValue });
          }}
        >
          {licenseTypes.map((type) => (
            <Option value={type}>{type}</Option>
          ))}
        </Select>
      </FormControl>
      <FormControl size="sm">
        <FormLabel>Architecture</FormLabel>
        <Select
          placeholder="All"
          disabled
          value={filters?.architecture}
          onChange={(e, newValue) => {
            setFilters({ ...filters, architecture: newValue });
          }}
        >
          {modelTypes.map((type) => (
            <Option value={type}>{type}</Option>
          ))}
        </Select>
      </FormControl>
    </>
  );

  return (
    <>
      <Box
        className="SearchAndFilters-tabletUp"
        sx={{
          borderRadius: 'sm',
          pb: 2,
          display: 'flex',
          flexWrap: 'wrap',
          gap: 1.5,
          '& > *': {
            minWidth: {
              xs: '120px',
              md: '160px',
            },
          },
        }}
      >
        <FormControl sx={{ flex: 1 }} size="sm">
          <FormLabel>&nbsp;</FormLabel>
          <Input
            placeholder="Search"
            value={searchText}
            onChange={(e) => setSearchText(e.target.value)}
            startDecorator={<SearchIcon />}
          />
        </FormControl>

        {renderFilters()}
      </Box>
      <Sheet
        className=""
        variant="outlined"
        sx={{
          width: '100%',
          borderRadius: 'md',
          minHeight: 0,
          display: 'flex',
          overflow: 'auto',
        }}
      >
        <Table
          aria-labelledby="tableTitle"
          stickyHeader
          hoverRow
          sx={{
            '--TableCell-headBackground': (theme) =>
              theme.vars.palette.background.level1,
            '--Table-headerUnderlineThickness': '1px',
            '--TableRow-hoverBackground': (theme) =>
              theme.vars.palette.background.level1,
            height: '100px',
            overflow: 'auto',
          }}
        >
          <thead>
            <tr>
              <th style={{ width: 140, padding: 12 }}>
                <Link
                  underline="none"
                  color="primary"
                  component="button"
                  onClick={() => setOrder(order === 'asc' ? 'desc' : 'asc')}
                  fontWeight="lg"
                  endDecorator={<ArrowDownIcon />}
                  sx={{
                    '& svg': {
                      transition: '0.2s',
                      transform:
                        order === 'desc' ? 'rotate(0deg)' : 'rotate(180deg)',
                    },
                  }}
                >
                  Name
                </Link>
              </th>
              <th style={{ width: 120, padding: 12 }}>Architecture</th>
              <th style={{ width: 120, padding: 12 }}>Params</th>
              {/* <th style={{ width: 220, padding: 12 }}>Type</th> */}
              <th style={{ width: 120, padding: 12 }}>Model ID</th>
              <th style={{ width: 160, padding: 12 }}> </th>
            </tr>
          </thead>
          <tbody>
            {models &&
              filterByFilters(models, searchText, filters).map((row) => (
                <tr key={row.rowid}>
                  <td>
                    <Typography ml={2} fontWeight="lg">
                      {row?.local_model === true ? (
                        <FlaskRoundIcon
                          color="var(--joy-palette-success-700)"
                          style={{
                            verticalAlign: 'middle',
                            marginRight: '5px',
                          }}
                        />
                      ) : (
                        ''
                      )}{' '}
                      {row.name}
                    </Typography>
                  </td>
                  <td>
                    {' '}
                    {row?.json_data?.architecture == 'MLX' && (
                      <>
                        <TinyMLXLogo />
                        &nbsp;
                      </>
                    )}
                    {row?.json_data?.architecture == 'GGUF' && (
                      <>
                        <img
                          src="https://avatars.githubusercontent.com/ggerganov"
                          width="24"
                          valign="middle"
                          style={{ borderRadius: '50%' }}
                        />{' '}
                        &nbsp;
                      </>
                    )}
                    {[
                      'LlamaForCausalLM',
                      'T5ForConditionalGeneration',
                      'FalconForCausalLM',
                      'MistralForCausalLM',
                      'GPTBigCodeForCausalLM',
                    ].includes(row?.json_data?.architecture) && <>🤗 &nbsp;</>}
                    {row?.json_data?.architecture}
                  </td>
                  <td>{row?.json_data?.parameters}</td>
                  {/* <td>{JSON.stringify(row)}</td> */}
                  {/* <td>
    <Box
      sx={{ display: "flex", gap: 2, alignItems: "center" }}
    ></Box>
  </td> */}
                  <td>{row.model_id}</td>
                  <td style={{ textAlign: 'right' }}>
                    {/* <Link fontWeight="lg" component="button" color="neutral">
    Archive
  </Link> */}
                    {pickAModelMode === true ? (
                      <SelectButton
                        setFoundation={setFoundation}
                        setAdaptor={setAdaptor}
                        model={row}
                      />
                    ) : (
                      <>
                        <InfoIcon
                          onClick={() => {
                            alert(JSON.stringify(row?.json_data));
                          }}
                        />
                        &nbsp;
                        <Trash2Icon
                          color="var(--joy-palette-danger-600)"
                          onClick={async () => {
                            if (
                              confirm(
                                "Are you sure you want to delete model '" +
                                  row.model_id +
                                  "'?"
                              )
                            ) {
                              await fetch(
                                chatAPI.Endpoints.Models.Delete(row.model_id)
                              );
                              mutateModels();
                            }
                          }}
                        />
                      </>
                    )}
                  </td>
                </tr>
              ))}
            {models?.length === 0 && (
              <tr>
                <td colSpan={5}>
                  <Typography
                    level="body-lg"
                    justifyContent="center"
                    margin={5}
                  >
                    You do not have any models on your local machine. You can
                    download a model by going to the{' '}
                    <ReactRouterLink to="/zoo">
                      <StoreIcon />
                      Model Store
                    </ReactRouterLink>
                    .
                  </Typography>
                </td>
              </tr>
            )}
          </tbody>
        </Table>
      </Sheet>
    </>
  );
};

export default LocalModelsTable;
